package com.xzx.flink.streamapi.state;

import com.xzx.flink.bean.ClickEvent;
import com.xzx.flink.streamapi.source.ClickSource;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;

import java.util.ArrayList;
import java.util.List;

/**
 * @version 1.0
 * @auther xinzhixuan
 * @date 2022/5/15 22:12
 */
public class OperatorState_01_Demo {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        SingleOutputStreamOperator<ClickEvent> stream = env.addSource(new ClickSource()).assignTimestampsAndWatermarks(WatermarkStrategy.<ClickEvent>forMonotonousTimestamps().withTimestampAssigner(new SerializableTimestampAssigner<ClickEvent>() {
            @Override
            public long extractTimestamp(ClickEvent element, long recordTimestamp) {
                return element.getTimestamp();
            }
        }));

        stream.print();

        stream.addSink(new BufferSink(10));


        env.execute(OperatorState_01_Demo.class.getSimpleName());
    }

    public static class BufferSink implements SinkFunction<ClickEvent>, CheckpointedFunction {

        private final int bufferSize;
        private final List<ClickEvent> buffer;
        private transient ListState<ClickEvent> listState;

        public BufferSink(int bufferSize) {
            this.bufferSize = bufferSize;
            buffer = new ArrayList<>();
        }

        @Override
        public void invoke(ClickEvent value, Context context) throws Exception {
            buffer.add(value);

            if (bufferSize == buffer.size()) {
                for (ClickEvent clickEvent : buffer) {
                    System.out.println(clickEvent);
                }

                System.out.println("=======批量输出完毕=====");
                buffer.clear();
            }
        }

        // checkpoint的时候被调用
        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            listState.clear();
            // 把当前局部变量中的所有元素写入到检查点中
            for (ClickEvent clickEvent : buffer) {
                listState.add(clickEvent);
            }
        }

        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            // getListState 说明使用的是轮训分布的方式平均分布
            listState = context.getOperatorStateStore().getListState(new ListStateDescriptor<ClickEvent>("buffer-elements", Types.POJO(ClickEvent.class)));
            // 如果是从故障中恢复，就将 ListState 中的所有元素添加到局部变量中
            if (context.isRestored()) {
                for (ClickEvent element : listState.get()) {
                    buffer.add(element);
                }
            }
        }
    }
}
